"""
文本主题分类: ClueAI/PromptCLUE，只开放了PromptCLUE-base版本，当前最新PromptCLUE-base-v1.5
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from threading import RLock
from transformers import pipeline
from typing import List

class TextClassifier(object):

    lock = RLock()
    def __new__(cls, *args, **kwargs):
        with TextClassifier.lock:
            if not hasattr(TextClassifier, "_instance"):
                TextClassifier._instance = object.__new__(cls)
                cls.pipe = pipeline("text-classification",
                                 model="classla/multilingual-IPTC-news-topic-classifier",
                                 device=0, max_length=512, truncation=True)
        return TextClassifier._instance

    def predict(self, texts: List[str]) -> List[str]:
        """
        sample：是否抽样。生成任务，可以设置为True;
        top_p：0-1之间，生成的内容越多样
        """
        res = self.pipe(texts)
        return [item['label'] for item in res]


if __name__ == '__main__':
    classifier = TextClassifier()
    print(classifier.predict(["如果日本沉没，中国会接收日本难民吗？"]))


